package com.xu.rag.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.date.DateUtil;
import cn.hutool.core.map.MapUtil;
import com.xu.rag.pojo.elasticData.ElasticVectorData;
import com.xu.rag.pojo.result.EmbeddingResult;
import com.xu.rag.service.VectorService;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.functionscore.ScriptScoreQueryBuilder;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.springframework.data.domain.Pageable;
import org.springframework.data.elasticsearch.core.*;
import org.springframework.data.elasticsearch.core.document.Document;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.IndexQuery;
import org.springframework.data.elasticsearch.core.query.IndexQueryBuilder;
import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.time.Instant;
import java.util.*;

/**
 * @Author: xdl
 * @CreateTime: 2025-01-19 21:59
 * @Description:
 */
@Slf4j
@Service
public class VectorServiceImpl implements VectorService {

    @Resource
    private ElasticsearchRestTemplate elasticsearchRestTemplate;

    public String getCollectionName() {
        return "llm_action_rag_" + DateUtil.format(Date.from(Instant.now()), "yyyyMMdd");
    }

    public boolean initCollection(String collectionName, int dim) {
        IndexOperations indexOperations = elasticsearchRestTemplate.indexOps(IndexCoordinates.of(collectionName));
        // 以当前日期算 当天的索引没创建就创建一下 方便测试玩
        if (!indexOperations.exists()) {
            Document document = Document.from(this.elasticMapping(dim));
            indexOperations.create(new HashMap<>(), document);
            return true;
        }
        return true;
    }

    public void storage(String collectionName, List<EmbeddingResult> embeddingResults) {
        List<IndexQuery> results = new ArrayList<>();
        for (EmbeddingResult embeddingResult : embeddingResults) {
            ElasticVectorData ele = new ElasticVectorData();
            ele.setVector(embeddingResult.getEmbedding());
            ele.setChunkId(embeddingResult.getRequestId());
            ele.setContent(embeddingResult.getPrompt());
            results.add(new IndexQueryBuilder().withObject(ele).build());
        }
        List<IndexedObjectInformation> bulkedResult = elasticsearchRestTemplate.bulkIndex(results, IndexCoordinates.of(collectionName));
        log.info("save vector success: {}", CollectionUtil.size(bulkedResult));
    }

    public String retrieval(String collectionName, double[] vector) {
        Map<String, Object> params = new HashMap<>();
        params.put("query_vector", vector);
        Script script = new Script(ScriptType.INLINE, Script.DEFAULT_SCRIPT_LANG, "cosineSimilarity(params.query_vector, 'vector')+1", params);
        ScriptScoreQueryBuilder scriptScoreQueryBuilder = new ScriptScoreQueryBuilder(QueryBuilders.boolQuery(), script);
        NativeSearchQuery nativeSearchQuery = new NativeSearchQueryBuilder()
                .withQuery(scriptScoreQueryBuilder)
                .withPageable(Pageable.ofSize(3))
                .build();
        SearchHits<ElasticVectorData> searchHits = this.elasticsearchRestTemplate.search(nativeSearchQuery, ElasticVectorData.class, IndexCoordinates.of(collectionName));
        List<SearchHit<ElasticVectorData>> data = searchHits.getSearchHits();
        List<String> results = new LinkedList<>();
        for (SearchHit<ElasticVectorData> ele : data) {
            results.add(ele.getContent().getContent());
        }
        return CollectionUtil.join(results, "");
    }

    private Map<String, Object> elasticMapping(int dims) {
        Map<String, Object> properties = new HashMap<>();
        properties
                .put("_class", MapUtil.builder("type", "keyword")
                .put("doc_values", "false")
                .put("index", "false")
                .build());
        properties
                .put("chunkId", MapUtil.builder("type", "keyword")
                .build());
        properties
                .put("content", MapUtil.builder("type", "keyword")
                .build());
        properties
                .put("docId", MapUtil.builder("type", "keyword")
                .build());
        properties
                .put("vector", MapUtil.builder("type", "dense_vector")
                .put("dims", Objects.toString(dims))
                .build());
        Map<String, Object> root = new HashMap<>();
        root.put("properties", properties);
        return root;
    }

}
